/*
 * @Author: chyang0822 270917365@qq.com
 * @Date: 2024-12-10 16:27:48
 * @LastEditors: chyang0822 270917365@qq.com
 * @LastEditTime: 2024-12-11 14:12:55
 * @FilePath: /tensorrt_guide/psenet/predict/layers.h
 * @Description: 
 * 
 * Copyright (c) 2024 by ${git_name_email}, All Rights Reserved. 
 */
#ifndef TENSORRTX_LAYERS_H
#define TENSORRTX_LAYERS_H

#include <map>
#include <math.h>
#include <assert.h>
#include <string>

#include "NvInfer.h"
#include "cuda_runtime_api.h"
using namespace nvinfer1;

IScaleLayer *addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, std::string lname, float eps);

IActivationLayer *bottleneck(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int ch, int stride, std::string lname, int branch_type);

IActivationLayer *addConvRelu(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int outch, int kernel, int stride, std::string lname);

#endif
